import os
import shutil
import time
import random
from selenium.webdriver.chrome.options import Options
from selenium.common.exceptions import  TimeoutException
from selenium.webdriver.common.by import By
from selenium.webdriver.support.ui import WebDriverWait
from selenium.webdriver.support import expected_conditions as ec
from selenium.webdriver.chrome.service import Service
from selenium import webdriver
import ddddocr
import sys

# 判断是否是打包环境
if getattr(sys, 'frozen', False):
    # 打包后的可执行文件所在目录
    BASE_DIR = os.path.dirname(sys.executable)
else:
    # 未打包时使用原逻辑
    BASE_DIR = os.path.dirname(os.path.abspath(__file__))

#存储最后下载的ID 和 登录账号密码
LAST_ID_FILE = os.path.join(BASE_DIR, "last_id.txt")
USER_INFO_FILE = os.path.join(BASE_DIR, "cnvd_cookie.txt")

# ChromeDriver路径
CHROMEDRIVER_PATH = os.path.join(BASE_DIR, "chromedriver")

# 下载文件夹
DOWNLOAD_DIR = BASE_DIR
TARGET_DIR = os.path.join(BASE_DIR,"../CNVD_DATA")

# CNVD基础URL模板
BASE_URL = "https://www.cnvd.org.cn/shareData/list"
#LOGIN_URL = "https://www.cnvd.org.cn/user/login"

# OCR工具
ocr = ddddocr.DdddOcr()

def get_last_id():
    if os.path.exists(LAST_ID_FILE):
        with open(LAST_ID_FILE, 'r', encoding='utf-8') as f:
            try:
                return int(f.read().strip())
            except:
                return 1
    else:
        return 1

def save_last_id(cid):
    with open(LAST_ID_FILE, 'w', encoding='utf-8') as f:
        f.write(str(cid))

def get_user_info():
    if os.path.exists(USER_INFO_FILE):
        with open(USER_INFO_FILE, 'r', encoding='utf-8') as f:
            lines = f.readlines()
            if len(lines) >= 2:
                email = lines[0].strip()
                password = lines[1].strip()
                return email, password
    return None, None

def delete_code_png():
    if os.path.exists('code_png.png'):
        os.remove('code_png.png')
        print("已删除 code_png.png")
    else:
        print("code_png.png 不存在")

def move_files_to_target_directory():
    if os.path.exists(DOWNLOAD_DIR):
        for filename in os.listdir(DOWNLOAD_DIR):
            if filename.endswith('.xml'):
                source_path = os.path.join(DOWNLOAD_DIR, filename)
                target_path = os.path.join(TARGET_DIR, filename)
                print(f"将xml文件从{DOWNLOAD_DIR}移动到{TARGET_DIR}")
                shutil.move(source_path, target_path)

# 截图并识别验证码
def recognize_captcha(driver):
    # 定位验证码图片元素
    try:
        captcha_img = WebDriverWait(driver, 10).until(
            ec.presence_of_element_located((By.ID, 'codeSpan'))
        )
        # 直接截图该元素
        captcha_img.screenshot('code_png.png')
    except Exception as e:
        print("验证码截图失败:", e)
        driver.save_screenshot('captcha_error.png')
        raise

    # OCR识别
    my_ocr = ddddocr.DdddOcr()
    with open('code_png.png', 'rb') as f:
        code_img_bytes = f.read()
    code = my_ocr.classification(code_img_bytes)

    # 切换标签页为登录页面
    driver.switch_to.window(driver.window_handles[0])
    print(f"识别的验证码: {code}")
    return code

# 登录逻辑
def login(driver, email, password):
    print("正在尝试登录...")
    try:
        email_input = WebDriverWait(driver, 10).until(ec.presence_of_element_located((By.ID, 'email')))
        password_input = driver.find_element(By.ID, 'password')
        code_input = driver.find_element(By.ID, 'myCode')
        time.sleep(2)
        email_input.clear()
        password_input.clear()
        code_input.clear()

        # 输入用户名、密码
        for attempt in range(3):  # 最多重试3次
            # 重新获取一下元素
            email_input = driver.find_element(By.ID, 'email')
            password_input = driver.find_element(By.ID, 'password')
            code_input = driver.find_element(By.ID, 'myCode')
            time.sleep(2)
            email_input.clear()
            password_input.clear()
            code_input.clear()

            for char in email:
                email_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))
            for char in password:
                password_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))
            time.sleep(2)  # 等待2秒后重试

            # 识别验证码
            code = recognize_captcha(driver)
            for char in code:
                code_input.send_keys(char)
                time.sleep(random.uniform(0.05, 0.15))

            # 提交登录
            login_button = driver.find_element(By.XPATH, '//a[@class="btn b_28_1"]')
            login_button.click()

            time.sleep(3)
            if "login" not in driver.current_url:
                print(" 登录成功！")
                return True

            print(f"登录失败，正在进行第{attempt + 1}次重试...")
            time.sleep(2)

    except Exception as e:
        print("登录过程中出错：", e)
        return False

def download_latest_file(driver, last_id):
    try:
        print("开始查找最新漏洞链接...")
        # 定位第一条记录的链接
        first_link = WebDriverWait(driver, 10).until(
            ec.element_to_be_clickable((By.XPATH, '//table[@class="tlist"]/tbody/tr[1]/td[1]/a'))
        )
        file_id = first_link.get_attribute('href').split("/")[-1]
        print(f"获取到漏洞ID: {file_id}")

        # 判断是否为新文件
        if int(file_id) <= last_id:
            print(f"当前最新漏洞ID {file_id}，小于等于上次已下载ID {last_id}，无需重复下载。")
            return last_id

        # 查找并点击下载链接
        download_link = WebDriverWait(driver, 10).until(
            ec.element_to_be_clickable((By.XPATH, '//a[@title="下载xml"]'))
        )

        # 再次确认下载链接对应的 ID 是否一致
        download_file_id = download_link.get_attribute('href').split("/")[-1]
        if download_file_id == file_id:
            print(f"下载链接: {download_link.get_attribute('href')}")
            download_link.click()
            print(f"成功触发下载漏洞文件: {file_id}.xml")
            # 休眠够时间，不然会受网络波动影响下载失败
            time.sleep(20)
            # 更新 last_id
            save_last_id(file_id)
            print(f"已更新 last_id 为: {file_id}")
            return file_id
        else:
            print("下载链接与当前漏洞ID不匹配，跳过下载。")
            return last_id

    except Exception as e:
        print("下载文件时发生错误：", e)
        return last_id

def main():

    if not os.path.exists(CHROMEDRIVER_PATH):
        print("chromedriver 未找到，请检查路径是否正确。")
        return

    # 浏览器设置
    chrome_options = Options()
    chrome_options.add_argument('--no-sandbox')
    chrome_options.add_argument('--disable-dev-shm-usage')
    chrome_options.add_argument('--disable-gpu')
    chrome_options.add_argument('--headless') # 注释之后可开启无头模式
    chrome_options.add_argument("--window-size=1920,1080")
    chrome_options.add_argument("user-agent=Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.0.0 Safari/537.36")
    # Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/103.0.9999.0 Safari/537.36
    chrome_options.add_argument("--disable-extensions")  # 禁用所有插件

    # 禁止自动化检测标志
    chrome_options.add_argument('--disable-blink-features=AutomationControlled')
    chrome_options.add_experimental_option("excludeSwitches", ["enable-automation"])
    chrome_options.add_experimental_option("useAutomationExtension", False)

    chrome_options.add_experimental_option("prefs", {
        "download.default_directory": DOWNLOAD_DIR,
        "download.prompt_for_download": False,
        "download.directory_upgrade": True,
        "safebrowsing.enabled": True
    })

    # 创建驱动
    chromedriver_path = CHROMEDRIVER_PATH
    service = Service(chromedriver_path)

    #driver = uc.Chrome(service=service, options=chrome_options)
    driver = webdriver.Chrome(service=service, options=chrome_options)

    try:
        current_last_id = get_last_id()
        print(f"当前记录的最后下载漏洞ID: {current_last_id}")

        # 获取缓存的账号密码
        email, password = get_user_info()
        if not email or not password:
            raise Exception("请查看cnnd_cookie.txt文件，修改账户密码信息")

        # 进入CNVD列表页
        print(f"正在访问URL: {BASE_URL}")
        driver.get(BASE_URL)

        # 检查是否需要登录
        try:
            WebDriverWait(driver, 5).until(
                ec.url_contains("login")
            )
            print("检测到登录页面，开始登录流程...")
            if not login(driver, email, password):
                raise Exception("登录失败，无法继续操作")
        except TimeoutException:
            print("当前已处于登录状态")

        # 开始下载
        new_last_id = download_latest_file(driver, current_last_id)
        print(f" 当前最新漏洞已下载完成，更新最后ID为: {new_last_id}")

        print(f"将文件移动到 {TARGET_DIR} 路径")
        move_files_to_target_directory()
    finally:
        # 添加短暂停顿，确保浏览器已关闭
        time.sleep(3)
        delete_code_png()
        driver.quit()

if __name__ == "__main__":
    main()

